import pymysql
from pymysql import cursors
from typing import Union, List, Dict, Any


class MySQLPlugin:
    """MySQL数据库操作插件

    功能特性：
    - 支持连接池管理
    - 自动提交事务机制
    - 类型提示和参数化查询
    - 灵活的查询结果返回格式
    - 上下文管理器支持

    示例配置：
    config = {
        'host': 'localhost',
        'port': 3306,
        'user': 'root',
        'password': 'secret',
        'database': 'my_db',
        'pool_size': 5
    }
    """

    def __init__(self, config: dict):
        self.config = config
        self.pool = []
        self._create_pool()

    def _create_pool(self):
        """创建连接池"""
        for _ in range(self.config.get('pool_size', 3)):
            conn = pymysql.connect(
                host=self.config['host'],
                port=self.config['port'],
                user=self.config['user'],
                password=self.config['password'],
                database=self.config['database'],
                cursorclass=cursors.DictCursor,
                autocommit=True
            )
            self.pool.append(conn)

    def get_conn(self):
        """从连接池获取连接"""
        if not self.pool:
            self._create_pool()
        return self.pool.pop()

    def release_conn(self, conn):
        """归还连接到连接池"""
        self.pool.append(conn)

    def close(self):
        """关闭连接池"""
        for conn in self.pool:
            conn.close()
        self.pool.clear()

    def execute(self, sql: str, params: Union[tuple, List[tuple]] = None) -> int:
        """执行写操作（INSERT/UPDATE/DELETE）

        Args:
            sql: SQL语句，使用%s作为占位符
            params: 单个元组（单条操作）或元组列表（批量操作）

        Returns:
            受影响的行数
        """
        conn = self.get_conn()
        try:
            with conn.cursor() as cursor:
                rowcount = cursor.execute(sql, params) if isinstance(params, tuple) \
                    else cursor.executemany(sql, params)
                return rowcount
        finally:
            self.release_conn(conn)

    def query(self, sql: str, params: tuple = None, fetch_one: bool = False) -> Union[List[dict], dict]:
        """执行查询操作

        Args:
            fetch_one: 是否只获取第一条结果
        """
        conn = self.get_conn()
        try:
            with conn.cursor() as cursor:
                cursor.execute(sql, params)
                return cursor.fetchone() if fetch_one else cursor.fetchall()
        finally:
            self.release_conn(conn)

    # 以下是快捷方法
    def insert(self, table: str, data: dict) -> int:
        """插入单条数据"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        return self.execute(sql, tuple(data.values()))

    def update(self, table: str, updates: dict, where: str, where_params: tuple) -> int:
        """更新数据"""
        set_clause = ', '.join([f"{k} = %s" for k in updates])
        sql = f"UPDATE {table} SET {set_clause} WHERE {where}"
        return self.execute(sql, tuple(updates.values()) + where_params)

    def delete(self, table: str, where: str, where_params: tuple) -> int:
        """删除数据"""
        sql = f"DELETE FROM {table} WHERE {where}"
        return self.execute(sql, where_params)

    # 事务支持
    def begin_transaction(self):
        """开启事务"""
        conn = self.get_conn()
        conn.autocommit(False)
        return conn

    def commit(self, conn):
        """提交事务"""
        conn.commit()
        self.release_conn(conn)

    def rollback(self, conn):
        """回滚事务"""
        conn.rollback()
        self.release_conn(conn)

    # 上下文管理器支持
    def __enter__(self):
        self.conn = self.begin_transaction()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.rollback(self.conn)
        else:
            self.commit(self.conn)


# 数据库相关内容填写
if __name__ == "__main__":
    DB_Host = ""
    DB_User = ""
    DB_Pwd = ""
    DB_Port = 3
    DB_Schema = ""
    config = {
        "host": DB_Host,
        "port": 3,
        "user": DB_User,
        "password": DB_Pwd,
        "database": DB_Schema,
        "pool_size": 3
    }

    db = MySQLPlugin(config)


